import asyncio
import aiohttp
import aiofiles
import ssl
import os
import concurrent.futures
from Crypto.Cipher import AES  # pip install pycryptodome
import requests
import time

# 在全局范围内禁用 SSL 证书验证
ssl_context = ssl.create_default_context()
ssl_context.check_hostname = False
ssl_context.verify_mode = ssl.CERT_NONE

# 设置全局线程池
executor = concurrent.futures.ThreadPoolExecutor()

# 下载m3u8文件
def download_m3u8_file(url, name):
    resp = requests.get(url)
    with open(f"实验数据/video2/m3u8文件/{name}", mode="wb") as f:
        f.write(resp.content)
    print('m3u8文件下载完成！')

# 根据文件提供数据下载ts
async def aio_download(qianzhui):
    tasks = []
    async with aiohttp.ClientSession() as session:  # 提前准备好session
        async with aiofiles.open(f"实验数据/video2/m3u8文件/m3u8.txt", mode="r", encoding="utf-8") as f:
            i = 0  # 计数器
            async for line in f:
                if line.startswith("#"):
                    continue
                # line就是xxxx.ts
                ts_url = qianzhui + line.strip()  # 去掉没用的空格和换行
                name = f"temp_{i}.ts"
                file_path = f"实验数据/video2/{name}"
                i += 1
                if not os.path.exists(file_path):
                    task = asyncio.create_task(download_ts(ts_url, name, session))  # 创建任务
                    tasks.append(task)

            if tasks:  # 检查是否有任务添加到列表中
                await asyncio.wait(tasks)  # 等待任务结束
            else:
                print("没有需要下载的文件！")
    print(f"{i}文件下载完毕！")

async def download_ts(url, name, session):
    try:
        async with session.get(url, ssl=ssl_context) as resp:
            async with aiofiles.open(f"实验数据/video2/{name}", mode="wb") as f:
                await f.write(await resp.content.read())  # 把下载到的内容写入到文件中
    except Exception as e:
        print(f"下载 {name} 时发生异常：{e}")

# 获得秘钥
def get_key(url):
    resp = requests.get(url)
    return resp.content

# 解密函数，处理一批次文件
async def aio_dec(key, start, end):
    tasks = []
    for i in range(start, end):
        task = asyncio.create_task(dec_ts(f"temp_{i}.ts", key))
        tasks.append(task)
    await asyncio.gather(*tasks)
    print(f"解密完毕: {start} 至 {end}")

# 解密单个文件
async def dec_ts(name, key):
    file_path = f"实验数据/video2/{name}"
    if not os.path.exists(file_path):
        print(f"文件不存在: {name}")
        return
    aes = AES.new(key=key, IV=b"0000000000000000", mode=AES.MODE_CBC)
    async with aiofiles.open(file_path, mode="rb") as f1, \
            aiofiles.open(f"实验数据/video2/po_{name}", mode="wb") as f2:
        bs = await f1.read()
        await f2.write(aes.decrypt(bs))

# 分批处理解密任务
async def batch_decrypt(key, num_files, max_batch_size):
    num_batches = (num_files + max_batch_size - 1) // max_batch_size
    for batch in range(num_batches):
        start = batch * max_batch_size
        end = min(start + max_batch_size, num_files)
        await aio_dec(key, start, end)

# 合并TS文件
def merge_ts(input_dir, output_file, max_files, ifjiami):
    with open(output_file, 'wb') as merged_file:
        for i in range(max_files):
            filename = f"{ifjiami}_{i}.ts"
            file_path = os.path.join(input_dir, filename)
            if not os.path.exists(file_path):
                print(f"文件不存在: {filename}")
                continue
            with open(file_path, 'rb') as source_file:
                merged_file.write(source_file.read())
    # 批量删除.ts文件
    ts_files = [f for f in os.listdir(input_dir) if f.endswith('.ts')]
    for ts_file in ts_files:
        os.unlink(os.path.join(input_dir, ts_file))
    print("完成")

# 主程序
def main():
    while True:
        print("请选择操作:")
        print("0. 下载m3u8文件")
        print("1. 下载批量文件")
        print("2. 解密文件")
        print("3. 合并视频")
        print("4. 退出")
        choice = input("输入选项(0/1/2/3/4): ")

        if choice == '0':
            m3u8_url = input("输入m3u8URL: ")
            try:
                download_m3u8_file(m3u8_url, "m3u8.txt")
            except RuntimeError as e:
                print(f"发生异常；{e}")

        elif choice == '1':
            qianzhui = input("输入下载链接前缀网址: ")
            print("开始下载批量文件...")
            try:
                loop = asyncio.get_event_loop()
                loop.set_default_executor(executor)
                loop.run_until_complete(aio_download(qianzhui))
            except RuntimeError as e:
                print(f"发生异常：{e}")
        elif choice == '2':
            key_url = input("请输入密钥URL: ")
            num_files = int(input("请输入需要解密的文件数: "))
            print("开始解密文件...")
            max_batch_size = 3000
            key = get_key(key_url)
            print("从URI获取的密钥:", key)
            print("密钥长度:", len(key))

            start_time = time.time()
            asyncio.run(batch_decrypt(key, num_files, max_batch_size))
            end_time = time.time()
            print(f"用时：{end_time - start_time} 秒")
        elif choice == '3':
            hebing = input("是不是合并加密文件？0不加密，1加密: ")
            if hebing == '0':
                ifjiami = "temp"
            else:
                ifjiami = "po_temp"
            input_directory = '实验数据/video2'  # 修改为包含.ts文件的目录
            output_filename = input("请输入合并后视频的名称: ")
            output_file = f'F:\\BaiduNetdiskDownload\\杂物\\{output_filename}.ts'
            max_files = int(input("请输入需要合并的文件数: "))
            print("开始合并视频...")
            merge_ts(input_directory, output_file, max_files, ifjiami)
        elif choice == '4':
            print("退出程序。")
            break
        else:
            print("无效选项，请重新选择。")

if __name__ == '__main__':
    main()
